import numpy as np
import pandas as pd
import os
from bishe_benchmark import get_all_nguyen_equations, generate_nguyen_data

def generate_and_save_data():
    # 创建数据保存目录
    data_dir = "nguyen_data"
    os.makedirs(data_dir, exist_ok=True)
    
    # 获取所有Nguyen方程
    equations = get_all_nguyen_equations()
    
    for func, name in equations:
        print(f"正在生成方程 {name} 的数据...")
        
        # 生成数据
        X, y = generate_nguyen_data(func, n_samples=1000, noise=0.1)
        
        # 创建DataFrame
        df = pd.DataFrame({
            'X': X.reshape(-1),
            'y': y.reshape(-1)
        })
        
        # 保存到CSV文件
        output_file = os.path.join(data_dir, f"{name}_data.csv")
        df.to_csv(output_file, index=False)
        print(f"数据已保存到: {output_file}")

if __name__ == "__main__":
    generate_and_save_data() 